import sys
import os
from typing import List, Optional
from datetime import datetime, timedelta
import google.generativeai as genai
from dotenv import load_dotenv
from .DataModels import ChatMessage, ChatRecord, OprationChain
import DataCollectionTools as DCTools
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
class ChatRecordProcessor:
    """聊天记录处理器"""
    async def get_formatted_chat_record(self, user_id: str, question_id: str, agent_type: str = "attack") -> ChatRecord:
        """
        获取格式化的聊天记录
        Args:
            user_id (str): 用户ID
            question_id (str): 问题ID
            agent_type (str): Agent类型 ("attack" or "defense")
        Returns:
            ChatRecord: 格式化的聊天记录
        """
        # 获取原始聊天记录
        raw_chat_data = DCTools.get_chat_record(
            params={"user_id": user_id, "question_id": question_id, "agent_type": agent_type},
            headers=self.headers
        )
        # 创建聊天记录对象
        chat_record = ChatRecord(
            user_id=user_id,
            question_id=question_id,
            agent_type=agent_type
        )
        # 处理聊天消息
        messages = []
        if raw_chat_data and isinstance(raw_chat_data, list):
            for msg_data in raw_chat_data:
                message = ChatMessage(
                    message_id=msg_data.get("message_id", ""),
                    user_id=msg_data.get("user_id", user_id),
                    question_id=msg_data.get("question_id", question_id),
                    message_type=msg_data.get("message_type", "user"),
                    agent_type=msg_data.get("agent_type", agent_type),
                    content=msg_data.get("content", ""),
                    timestamp=msg_data.get("timestamp", ""),
                    # is_helpful=msg_data.get("is_helpful"),
                    # attack_action=msg_data.get("attack_action"),
                    # defense_action=msg_data.get("defense_action")
                )
                messages.append(message)
        # 按时间戳排序消息
        messages.sort(key=lambda x: x.timestamp)
        chat_record.messages = messages
        chat_record.total_messages = len(messages)
        # 计算对话持续时间
        if len(messages) >= 2:
            try:
                start_time = datetime.fromisoformat(messages[0].timestamp.replace('Z', '+00:00'))
                end_time = datetime.fromisoformat(messages[-1].timestamp.replace('Z', '+00:00'))
                chat_record.conversation_duration = end_time - start_time
            except Exception as e:
                print(f"计算对话持续时间失败: {e}")
                chat_record.conversation_duration = None
        chat_record.opration_chain = await self.get_opration_chain(chat_record)
        # 计算用户满意度（基于有用性评分）
        # helpful_messages = [msg for msg in messages if msg.is_helpful is not None]
        # if helpful_messages:
        #     helpful_count = sum(1 for msg in helpful_messages if msg.is_helpful)
        #     chat_record.user_satisfaction = helpful_count / len(helpful_messages)
        return chat_record
    async def get_opration_chain(self, chat_record: ChatRecord):
        """
        获取操作链
        Args:
            chat_record (ChatRecord): 聊天记录对象
        Returns:
            OprationChain: 操作链数据
        """
        # 创建操作链对象
        opration_chain = OprationChain(
            user_id=chat_record.user_id,
            question_id=chat_record.question_id
        )
        # 提取操作行为
        current_dir = os.path.dirname(__file__)
        project_root = os.path.abspath(os.path.join(current_dir, '..'))
        env_path = os.path.join(project_root, ".env")
        if not os.path.exists(env_path):
            print(f"Warning: .env file not found at {env_path}")
            return {} 
        load_dotenv(env_path)
        PROXY = os.getenv("PROXY")
        if PROXY:
            os.environ["http_proxy"] = PROXY
            os.environ["https_proxy"] = PROXY 
            os.environ["NO_PROXY"] = "127.0.0.1,localhost"
        API_KEY = os.getenv("GEMINI_API_KEY")
        if not API_KEY:
            print("Warning: GEMINI_API_KEY not found in environment variables")
            return {}
        genai.configure(api_key=API_KEY)
        prompt = """
            this is a chat record between user and attack or defense agent,
            you need to extract the opration(attack or defense) chain from the chat record,
            return the opration chain in a python str list format,
            for example:
            ["action1", "action2", "action3"]
        """
        messages = [
            {"role": "professional network security expert", "content": "you are a professional network security expert."},
            {"role": "user", "content": prompt + "\n\nchat record:\n" + "\n".join(msg.content for msg in chat_record.messages)}
        ]
        model = genai.GenerativeModel("gemini-2.5-flash")
        chat = model.start_chat()
        response = chat.send_message(str(messages))
        #提取回答列表
        response_list = response.split("\n")
        #提取操作链
        opration_chain = []
        for line in response_list:
            line = line.strip()
            if line.startswith("[") and line.endswith("]"):
                try:
                    opration_chain = eval(line)
                    if isinstance(opration_chain, list):
                        break
                except Exception as e:
                    print(f"Error parsing opration chain: {e}")
        return opration_chain
    def format_chat_for_display(self, chat_record: ChatRecord) -> str:
        """
        将聊天记录格式化为易读的显示格式
        Args:
            chat_record (ChatRecord): 聊天记录对象
        Returns:
            str: 格式化的聊天记录字符串
        """
        lines = []
        agent_type_display = "攻击Agent" if chat_record.agent_type == "attack" else "防御Agent"
        lines.append(f"=== {agent_type_display}聊天记录 ===")
        lines.append(f"用户ID: {chat_record.user_id}")
        lines.append(f"问题ID: {chat_record.question_id}")
        lines.append(f"Agent类型: {agent_type_display}")
        lines.append(f"消息总数: {chat_record.total_messages}")
        if chat_record.conversation_duration:
            lines.append(f"对话持续时间: {chat_record.conversation_duration}")
        # if chat_record.user_satisfaction is not None:
        #     lines.append(f"用户满意度: {chat_record.user_satisfaction:.2%}")
        lines.append("")
        lines.append("=== 对话内容 ===")
        for message in chat_record.messages:
            role = "🧑 用户" if message.message_type == "user" else "🤖 助手"
            timestamp = message.timestamp[:19] if message.timestamp else "未知时间"
            lines.append(f"[{timestamp}] {role}:")
            lines.append(f"  {message.content}")
            # 显示攻击行为（仅攻击Agent消息）
            if message.attack_action and chat_record.agent_type == "attack":
                lines.append(f"  ⚔️ 攻击行为: {message.attack_action}")
            # 显示防御措施（仅防御Agent消息）
            if message.defense_action and chat_record.agent_type == "defense":
                lines.append(f"  🛡️ 防御措施: {message.defense_action}")
            if message.is_helpful is not None:
                helpful_text = "✅ 有用" if message.is_helpful else "❌ 无用"
                lines.append(f"  ({helpful_text})")
            lines.append("")
        return "\n".join(lines)
    def get_chat_statistics(self, chat_record: ChatRecord) -> dict:
        """
        获取聊天记录的统计信息
        Args:
            chat_record (ChatRecord): 聊天记录对象
        Returns:
            dict: 统计信息
        """
        user_messages = [msg for msg in chat_record.messages if msg.message_type == "user"]
        assistant_messages = [msg for msg in chat_record.messages if msg.message_type == "assistant"]
        stats = {
            "total_messages": chat_record.total_messages,
            "user_messages": len(user_messages),
            "assistant_messages": len(assistant_messages),
            "conversation_duration": chat_record.conversation_duration,
            "user_satisfaction": chat_record.user_satisfaction,
            "avg_message_length": {
                "user": sum(len(msg.content) for msg in user_messages) / len(user_messages) if user_messages else 0,
                "assistant": sum(len(msg.content) for msg in assistant_messages) / len(assistant_messages) if assistant_messages else 0
            }
        }
        return stats
